import { llm } from "../utils/llm";
import { Annotation } from "@langchain/langgraph";
import { write, Xlsx } from "../utils/xlsx";
import { OutputState } from "../utils/state";
import { log } from "../utils/logs";
import { prompt } from "../utils/prompt";
import { _parse } from "../utils/markdown";

export const CargoState = Annotation.Root({
  file: Annotation<string>(), // 文件
  xlsx: Annotation<Xlsx>(), // 读取的xlsx文件

  table: Annotation<any>(), // 货物清单

  ...OutputState.spec,
});

/**
 * 货物清单
 */
export const _cargo = async (state: typeof CargoState.State) => {
  const { file, xlsx } = state;

  await log(file, "format.md", "");

  const recursive = async (input: any) => {
    let markdown = "";
    const chunks: any = await llm({
      model: "qwen-plus",
      stream: true,
      messages: [
        {
          role: "system",
          content: await prompt("_cargo"),
        },
        {
          role: "user",
          content: JSON.stringify(input),
        },
      ],
    });

    for await (const chunk of chunks) {
      const content = chunk.choices?.[0]?.delta?.content;
      markdown += content ?? "";
      await log(file, "format.md", content);
    }
    await log(file, "format.md", "\n\n");
    return markdown;
  };

  let markdown = "";
  const document = xlsx.document ?? [];
  for (let i = 0; i < document.length; i++) {
    const sheet = document?.[i];
    const sheetName = sheet?.name;
    const sheetData = sheet?.data;
    if (!sheetName || !sheetData) {
      continue;
    }
    let b = `### ${sheetName}\n`;
    for (let j = 0; j < sheetData?.length; j += 10) {
      const latest =
        j + 10 > sheetData?.length ? sheetData?.length - 1 : j + 10;
      const d = sheetData?.slice(j, latest);
      b += `\n`;
      b += await recursive(d);
      b += `\n`;
    }
    b += "\n\n";
    markdown += b;
  }

  const ast = await _parse(markdown);
  const records = ast?.children
    ?.filter((node: any) => node.type === "code" && "json" === node.lang)
    ?.map((node: any) => JSON.parse(node.value ?? "[]"))
    ?.filter((row: any) => row)
    ?.reduce((prev: any, curr: any) => prev.concat(curr), []);
  await log(file, "format.json", JSON.stringify(records, null, 2));
  return {
    ...state,
    records,
  };
};

/**
 * 货物清单
 */
export const __cargo = async (state: typeof CargoState.State) => {
  const { file, xlsx } = state;

  await log(file, "format.md", "");

  // 预测，获取表头和开始行
  const preview = async (sheetData: any) => {
    let markdown = "";
    const chunks: any = await llm({
      model: "qwen-plus",
      stream: true,
      messages: [
        {
          role: "system",
          content: await prompt("preview"),
        },
        {
          role: "user",
          content: sheetData?.map((row: any[]) => row.join(", "))?.join("\n"),
        },
      ],
    });

    for await (const chunk of chunks) {
      const content = chunk.choices?.[0]?.delta?.content;
      markdown += content ?? "";
      await log(file, "preview.md", content, true);
    }
    await log(file, "preview.md", "\n\n", true);
    return markdown;
  };

  // 轮询，解决token长度问题
  const recursive = async (input: any, headers?: any) => {
    let markdown = "";
    const chunks: any = await llm({
      model: "qwen-plus",
      stream: true,
      messages: [
        {
          role: "system",
          content: await prompt("cargo", {
            headers: headers ?? {},
          }),
        },
        {
          role: "user",
          content: input?.map((row: any[]) => row.join(", "))?.join("\n"),
        },
      ],
    });

    for await (const chunk of chunks) {
      const content = chunk.choices?.[0]?.delta?.content;
      markdown += content ?? "";
      await log(file, "format.md", content, true);
    }
    await log(file, "format.md", "\n\n", true);
    return markdown;
  };

  // 按指定格式提取内容
  const format = async (sheetData: any) => {
    let markdown = await recursive(sheetData);
    const ast = await _parse(markdown);
    const info: any = ast?.children?.filter(
      (node) =>
        "code" === node?.type && "json" === node?.lang && "info" == node?.meta
    )?.[0];
    if (info) {
      try {
        const headers = JSON.parse(info?.value ?? "{}");
        const latest = headers?.latest
          ?.map((it: any) => it?.toString() ?? "null")
          .json(",");
        const index = sheetData?.findLastIndex((row: any) => {
          const r = row?.map((it: any) => it?.toString() ?? "null").json(",");
          return latest === r;
        });
        const from = index + 1;
        const next = sheetData?.slice(from);
        if (next && next?.length > 0) {
          markdown += `\n\n------\n\n`;
          markdown += await format(next);
        }
      } catch (error) {
        console.log("继续提取内容异常", error);
      }
    }
    return markdown;
  };

  // 输入表格，提取内容
  let markdown = "";
  const document = xlsx.document ?? [];
  for (let i = 0; i < document.length; i++) {
    const sheet = document?.[i];
    const sheetName = sheet?.name;
    const sheetData = sheet?.data;
    if (!sheetName || !sheetData) {
      continue;
    }
    // await preview(sheetData);
    let b = `### ${sheetName}\n`;
    b += await format(sheetData);
    b += "\n\n";
    markdown += b;
  }

  // 容错处理，解析，获取目标结果
  const ast = await _parse(markdown);
  const records = ast?.children
    ?.filter(
      (node: any) =>
        "code" === node?.type &&
        "json" === node?.lang &&
        "content" === node?.meta
    )
    ?.map((node: any) => {
      console.log(node);
      return node && JSON.parse(node.value ?? "[]");
    })
    ?.filter((row: any) => row)
    ?.reduce((prev: any, curr: any) => prev.concat(curr), []);
  await log(file, "format.json", JSON.stringify(records, null, 2));
  return {
    ...state,
    records,
  };
};

/**
 * 货物清单
 */
export const cargo = async (state: typeof CargoState.State) => {
  const { file, xlsx } = state;

  await log(file, "format.md", "");

  // 获取表头和开始行
  const headers = async (sheetName: string, sheetData: any) => {
    let markdown = "";
    const chunks: any = await llm({
      model: "qwen-plus",
      stream: true,
      messages: [
        {
          role: "system",
          content: await prompt("headers"),
        },
        {
          role: "user",
          content: sheetData?.map((row: any[]) => row.join(", "))?.join("\n"),
        },
      ],
    });

    for await (const chunk of chunks) {
      const content = chunk.choices?.[0]?.delta?.content;
      markdown += content ?? "";
      await log(file, `${sheetName}headers.md`, content, true);
    }
    await log(file, `${sheetName}headers.md`, "\n\n", true);
    const ast = await _parse(markdown);
    const node: any = ast?.children?.filter((n: any) => {
      return "code" === n?.type && "json" === n?.lang && "head" === n?.meta;
    })?.[0];
    let head = {
      headers: [],
    };
    try {
      head = JSON.parse(node?.value ?? "{}");
    } catch (err) {
      console.log("解析表头异常", err);
    }
    return head;
  };

  // 轮询，解决token长度问题
  const recursive = async (input: any, head?: any) => {
    let markdown = "";
    const chunks: any = await llm({
      model: "qwen-plus",
      stream: true,
      messages: [
        {
          role: "system",
          content: await prompt("cargo"),
        },
        {
          role: "user",
          content: await prompt("input", {
            head,
            data: input?.map((it: any[]) => it.join("\t")).join("\n"),
          }),
        },
      ],
    });

    for await (const chunk of chunks) {
      const content = chunk.choices?.[0]?.delta?.content;
      markdown += content ?? "";
      await log(file, "format.md", content, true);
    }
    await log(file, "format.md", "\n\n", true);
    return markdown;
  };

  // 输入表格，提取内容
  let markdown = "";
  let table = {
    file,
    document: [],
  };
  const document = xlsx.document ?? [];
  for (let i = 0; i < document.length; i++) {
    const sheet = document?.[i];
    const sheetName = sheet?.name;
    const sheetData = sheet?.data;
    if (!sheetName || !sheetData) {
      continue;
    }
    let head: any = await headers(sheetName, sheetData);
    const doc = {
      name: sheetName,
      data: [],
    };
    let b = `### ${sheetName}\n\n`;
    await log(file, "format.md", `### ${sheetName}\n`, true);

    // 每10条数据进行一次处理
    for (let j = 0; j < sheetData?.length; j += 10) {
      const next = j + 10 > sheetData?.length ? sheetData?.length : j + 10;
      const slicedData = sheetData?.slice(j, next);
      const res = await recursive(slicedData, head);
      b += `${res}\n`;

      try {
        const ast = await _parse(res);
        const list = ast?.children
          ?.filter(
            (node: any) =>
              "code" === node?.type &&
              "json" === node?.lang &&
              "content" === node?.meta
          )
          ?.map((node: any) => {
            return node && JSON.parse(node.value ?? "[]");
          })
          ?.filter((row: any) => row)
          ?.reduce((prev: any, curr: any) => prev.concat(curr), []);
        (doc.data as any[]).push(...list);
      } catch (err) {
        console.error("解析内容异常", err);
      }

      try {
        const node: any = (await _parse(res))?.children?.filter((n: any) => {
          return (
            "code" === n?.type && "json" === n?.lang && "content" === n?.meta
          );
        })?.[0];
        const content = JSON.parse(node?.value ?? "[]");
        head = {
          ...head,
          latest: content.length > 0 ? content?.[content.length - 1] : {},
        };
      } catch (err) {
        console.error("解析表头异常", err);
      }
    }
    await log(file, `${sheetName}.json`, JSON.stringify(doc, null, 2));
    (table?.document as any[])?.push(doc);

    // 将doc写入xlsx文件
    const _h = [
      "货物名",
      "单件重量（单位：kg）",
      "件数",
      "总重量（单位：kg）",
      "长（单位: m）",
      "宽（单位: m）",
      "高（单位: m）",
      "体积（单位: m3）",
      "总体积（单位: m3）",
      "备注",
    ];
    const items: any[][] = doc?.data?.map((item) =>
      _h.map((key) => item[key] ?? "")
    );
    await write(file, sheetName, _h, items);

    await log(file, "format.md", `\n`, true);
    b += "\n";
    markdown += b;
  }

  await log(file, "data.md", markdown);
  await log(file, "table.json", JSON.stringify(table, null, 2));
  return {
    ...state,
    table,
  };
};
